
from managers.nodes.baseGraphNode import *
from agent.common_base_agent import CommonBaseAgent
from agent.task_list_split_agent import TaskListSplitAgent
from utils.msgUtils import take_text
from .sessionInfo import SessionInfo
from ..execute.baseExeHis import BaseExeHis
from eventlet.semaphore import Semaphore
from collections import defaultdict


class AgentTaskListGraphNode(BaseGraphNode):
    def __init__(self, id:str, data:dict, graphNodeManager, random_key:str):
        super().__init__(id, data, graphNodeManager, random_key)
        self.name = data['data']['name']
        self.role = data['data'].get('role', 'assistant')
        self.prompt = data['data'].get('prompt', 'You are a helpful assistant.')
        self.subrole = data['data'].get('subrole', 'assistant')
        self.subprompt = data['data'].get('subprompt', 'You are a helpful assistant.')
        self.subusingcontext = data['data'].get('subusingcontext', True)
        self.subusingmemory = data['data'].get('subusingmemory', True)
        self.subisSpeak = data['data'].get('subisSpeak', True)
        self.subisFunctionCall = data['data'].get('subisFunctionCall', False)
        self.customKey = None
        if len(data['data'].get('custom_keys', []))>=1:
            self.customKey = data['data']['custom_keys'] # [{"class_type": "AgentMsg", "handle_type": "source", "id": "custom-key-total_score-msg", "name": "total_score", "note": "\u603b\u5206", "top": 65, "type": "str", "value_type": "msg"}]
        self.uuid2agent = {}
        self.session_locks = defaultdict(Semaphore)  # 每个session_id独立锁
        

    def execute(self, session_info: SessionInfo=None, exe_his: BaseExeHis=None, trigger=None, running_id=None):
        try:
            if session_info is None: return 
            query = Msg("","user","text")
            if session_info.id not in self.uuid2agent:
                self.uuid2agent[session_info.id] = TaskListSplitAgent(name = self.role, sys_prompt = self.prompt, model_config_name = 'openai_cfg', isSpeak = False, isFunctionCall = False)
            agent = self.uuid2agent[session_info.id]
            if 'query-msg' in self.dict['in']: # {'nodeId': sourceID, 'handle': sourceHandle}
                queryFrom = self.dict['in']['query-msg']
                prenode =self.graphNodeManager.nodesByIdDict[queryFrom['nodeId']]
                query = prenode.get_newest_output_in_session(session_info, queryFrom['handle'])
                if query is None:
                    query = Msg("The user has not provided any query. Refer to your system prompt.", "user", "text")
            
            if 'context-msg' in self.dict['in']:
                contextFrom = self.dict['in']['context-msg']
                prenode = self.graphNodeManager.nodesByIdDict[contextFrom['nodeId']]
                context = prenode.get_newest_output_in_session(session_info, contextFrom['handle'])
                if context is not None:
                    agent.context = context.content

            if 'memory-msgs' in self.dict['in']:
                memoryFrom = self.dict['in']['memory-msgs']
                prenode = self.graphNodeManager.nodesByIdDict[memoryFrom['nodeId']]
                memorys = prenode.get_newest_output_in_session(session_info, memoryFrom['handle'])
                agent.setmemory(memorys)

            if 'tools' in self.dict['in']:
                toolsFrom = self.dict['in']['tools']
                prenode = self.graphNodeManager.nodesByIdDict[toolsFrom['nodeId']]
                tools = prenode.get_newest_output_in_session(session_info, toolsFrom['handle'])
                agent.tools = tools
        
            session_info.ws.start_background_task(self._process_agent_reply, session_info, query, exe_his, trigger, running_id)
        except Exception as e:
            print("agentGraphNode execute error", e)
            return None        


    def _process_agent_reply(self, session_info: SessionInfo, query, exe_his, trigger, running_id):
        returnMsg = self.uuid2agent[session_info.id].reply(query.msg_dict(), exe_his)
        for query in returnMsg['querys']:
            if query is None: continue
            task_temp_agent = CommonBaseAgent(name = self.subrole, sys_prompt = self.subprompt, model_config_name = 'openai_cfg', isSpeak = self.subisSpeak, isFunctionCall = self.subisFunctionCall, customKey = self.customKey)
            if self.subusingcontext:
                task_temp_agent.context = self.uuid2agent[session_info.id].context
            if self.subusingmemory:
                task_temp_agent.setmemory(self.uuid2agent[session_info.id].memory)
            task_temp_agent.tools = self.uuid2agent[session_info.id].tools
            subreturnMsg = task_temp_agent.reply(Msg(query, 'user', 'text').msg_dict(), exe_his)
            with session_info.out_write_lock:
                if session_info.out.get(self.id,None) is None: session_info.out[self.id]={}
                if('reason' in subreturnMsg):
                    if 'reason-msg' not in session_info.out[self.id]: session_info.out[self.id]['reason-msg'] = []
                    session_info.out[self.id]['reason-msg'].append(Msg(subreturnMsg['reason'],'assistant','text')  )
                if('speak' in subreturnMsg):
                    if 'speak-msg' not in session_info.out[self.id]: session_info.out[self.id]['speak-msg'] = []
                    session_info.out[self.id]['speak-msg'].append(Msg(subreturnMsg['speak'],'assistant','text') )
                if ('functions' in subreturnMsg):
                    if('function-calls' not in session_info.out[self.id]): session_info.out[self.id]['function-calls'] = []
                    session_info.out[self.id]['function-calls'].append(subreturnMsg['functions'])
                if self.customKey:
                    for ck in self.customKey:
                        if ck['id'] not in session_info.out[self.id]: session_info.out[self.id][ck['id']] = []
                        session_info.out[self.id][ck['id']].append(  subreturnMsg[ck['name']]  )
                self.executed_num += 1
                if exe_his:exe_his.done()
                if trigger and trigger.enInterrupt and trigger.session_id_2_running_id[session_info.id] != running_id: return 
            self.flow_next(session_info, trigger=trigger, running_id=running_id, type='text', role='assistant')
    def flow_next(self, session_info: SessionInfo, **kwargs):
        return super().flow_next(session_info, **kwargs)

    def clear(self, session_info: SessionInfo):
        with session_info.out_write_lock:
            if session_info.out.get(self.id,None) is None: session_info.out[self.id]={}
        session_info.out[self.id].clear()
        if session_info.id in self.uuid2agent:
            self.uuid2agent[session_info.id].clear()
